package com.example.speechtotext.demos.web.handler;

import com.alibaba.nls.client.protocol.InputFormatEnum;
import com.alibaba.nls.client.protocol.NlsClient;
import com.alibaba.nls.client.protocol.SampleRateEnum;
import com.alibaba.nls.client.protocol.asr.SpeechTranscriber;
import com.alibaba.nls.client.protocol.asr.SpeechTranscriberListener;
import com.alibaba.nls.client.protocol.asr.SpeechTranscriberResponse;
import jakarta.annotation.PostConstruct;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.BinaryWebSocketHandler;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @description:
 * @author: MS
 * @date: 2025/3/11 16:40
 */
@Component
public class SpeechWebSocketHandler extends BinaryWebSocketHandler {

    private static final Logger logger = LoggerFactory.getLogger(SpeechWebSocketHandler.class);

    // 应用key，需要使用自己的应用key
    private String appKey = "D7WYgzqHTX7Qtvn7";
    // 临时token，需要重新申请
    private String accesToken = "32e3167cce8f42d8af630cf26edf052f";
    private static final int MAX_CONCURRENT = 100;


    private NlsClient nlsClient;
    private final ConcurrentHashMap<String, SpeechTranscriber> transcriberMap = new ConcurrentHashMap<>();

    @PostConstruct
    public void init() throws IOException {
        // 初始化阿里云客户端（应用启动时创建一次）
        this.nlsClient = new NlsClient(accesToken);
    }

    @Override
    public void afterConnectionEstablished(@NotNull WebSocketSession session) throws Exception {
        // 超出最大并发数时，拒绝连接
        if (transcriberMap.size() >= MAX_CONCURRENT) {
            session.close(new CloseStatus(429, "Too Many Requests"));
            return;
        }

        // 每个会话创建独立的SpeechTranscriber（用来向阿里发送数据）
        SpeechTranscriber transcriber = createTranscriber(session);
        transcriberMap.put(session.getId(), transcriber);
        session.sendMessage(new TextMessage("-------CONNECTED-------"));
    }

    // 创建SpeechTranscriber，并绑定监听器（listener）和session
    private SpeechTranscriber createTranscriber(WebSocketSession session) throws Exception {
        // 设置监听器，监听识别结果
        SpeechTranscriberListener listener = new SpeechTranscriberListener() {
            // 会话启动
            @Override
            public void onTranscriberStart(SpeechTranscriberResponse response) {
                String onTranscriberStart = "onTranscriberStart:\n" +
                        "taskId: " + response.getTaskId() +
                        ", status: " + response.getStatus() +
                        "（20000000=正常，其他为异常）";
                if (session.isOpen()) {
                    try {
                        session.sendMessage(new TextMessage(onTranscriberStart));
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                    System.out.println(onTranscriberStart);
                }
            }

            // 句子开始
            @Override
            public void onSentenceBegin(SpeechTranscriberResponse response) {
                String onSentenceBeginStr = "onSentenceBegin: " + response.getTransSentenceText();
                if (session.isOpen()) {
                    try {
                        session.sendMessage(new TextMessage(onSentenceBeginStr));
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                    System.out.println(onSentenceBeginStr);
                }
            }

            // 中间字段
            @Override
            public void onTranscriptionResultChange(SpeechTranscriberResponse response) {
                String onTranscriptionResultChangeStr = "onTranscriptionResultChange:\n" +
                        "句子编号: " + response.getTransSentenceIndex() +
                        ", 临时文本: \"" + response.getTransSentenceText() + "\"" +
                        ", 已处理: " + response.getTransSentenceTime() + "ms";
                if (session.isOpen()) {
                    try {
                        session.sendMessage(new TextMessage(onTranscriptionResultChangeStr));
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                    System.out.println(onTranscriptionResultChangeStr);
                }
            }

            // 句子结束
            @Override
            public void onSentenceEnd(SpeechTranscriberResponse response) {
                String onSentenceEndStr = "onSentenceEnd: " + response.getTransSentenceText();
                if (session.isOpen()) {
                    try {
                        session.sendMessage(new TextMessage(onSentenceEndStr));
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                    System.out.println(onSentenceEndStr);
                }
            }

            // 识别结束
            @Override
            public void onTranscriptionComplete(SpeechTranscriberResponse response) {
                String onTranscriptionCompleteStr = "onTranscriptionComplete: " + response.getStashResult();
                if (session.isOpen()) {
                    try {
                        session.sendMessage(new TextMessage(onTranscriptionCompleteStr));
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                    System.out.println(onTranscriptionCompleteStr);
                }
            }

            // 识别失败
            @Override
            public void onFail(SpeechTranscriberResponse response) {
                String onFailStr = "onFail: " + response.getTaskId() +
                        ", status: " + response.getStatus() +
                        ", status_text: " + response.getStatusText();
                if (session.isOpen()) {
                    try {
                        session.sendMessage(new TextMessage(onFailStr));
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                    System.out.println(onFailStr);
                }
            }
        };

        SpeechTranscriber transcriber = new SpeechTranscriber(nlsClient, listener);
        transcriber.setAppKey(appKey);
        transcriber.setFormat(InputFormatEnum.PCM);
        transcriber.setSampleRate(SampleRateEnum.SAMPLE_RATE_16K);
        //是否返回中间识别结果。
        transcriber.setEnableIntermediateResult(true);
        //是否生成并返回标点符号。
        transcriber.setEnablePunctuation(true);
        //是否将返回结果规整化，比如将一百返回为100。
        transcriber.setEnableITN(false);

        //设置vad断句参数。默认值：800ms，有效值：200ms～6000ms。
        transcriber.addCustomedParam("max_sentence_silence", 600);
        //设置是否语义断句。
        transcriber.addCustomedParam("enable_semantic_sentence_detection",false);
        //设置是否开启过滤语气词，即声音顺滑。
        transcriber.addCustomedParam("disfluency",true);
        //设置是否开启词模式。
        transcriber.addCustomedParam("enable_words",true);
        //设置vad噪音阈值参数，参数取值为-1～+1，如-0.9、-0.8、0.2、0.9。
        //取值越趋于-1，判定为语音的概率越大，亦即有可能更多噪声被当成语音被误识别。
        //取值越趋于+1，判定为噪音的越多，亦即有可能更多语音段被当成噪音被拒绝识别。
        //该参数属高级参数，调整需慎重和重点测试。
        //transcriber.addCustomedParam("speech_noise_threshold",0.3);
        //设置训练后的定制语言模型id。
        //transcriber.addCustomedParam("customization_id","你的定制语言模型id");
        //设置训练后的定制热词id。
        //transcriber.addCustomedParam("vocabulary_id","你的定制热词id");

        //此方法将以上参数设置序列化为JSON发送给服务端，并等待服务端确认。
        transcriber.start();
        return transcriber;
    }

    @Override
    protected void handleBinaryMessage(@NotNull WebSocketSession session, @NotNull BinaryMessage message) throws IOException {
        String sessionId = session.getId();
        SpeechTranscriber transcriber = transcriberMap.get(sessionId);
        if (transcriber == null) {
            logger.error("会话 {} 的识别器未初始化", sessionId);
            return;
        }

        // 接受前端发送的音频数据，将数据发送到阿里云服务器进行识别
        ByteBuffer buffer = message.getPayload();
        byte[] chunk = new byte[1024]; // 16K采样率下，每20ms约320字节（16bit*16000Hz*0.02s）
        while (buffer.hasRemaining()) {
            int len = Math.min(buffer.remaining(), chunk.length);
            buffer.get(chunk, 0, len);
            if (session.isOpen()) {
                transcriber.send(chunk); // 分块发送，降低内存压力
            }
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, @NotNull CloseStatus status) throws Exception {
        String sessionId = session.getId();
        SpeechTranscriber transcriber = transcriberMap.remove(sessionId);
        if (transcriber != null) {
            transcriber.stop();
            transcriber.close();
            logger.info("释放识别器资源: {}", sessionId);
        }
    }
}